#! /usr/bin/env python3

import argparse
from glob import glob
from netCDF4 import Dataset 
import os
import re
import pendulum
from jinja2 import Template
import f90nml
import sys
script_root = os.path.dirname(os.path.realpath(__file__))
sys.path.append(f'{script_root}/../utils')
from utils import cli, check_files, run, parse_config, submit_job, has_key, get_value

def run_wrfda_obsproc(work_root, wrfda_root, config, args, littler_root=None, wrf_work_dir=None, tag=None, blacklist=False, whitelist=False):
	start_time = config['custom']['start_time']
	datetime_fmt = 'YYYY-MM-DD_HH:mm:ss'
	start_time_str = start_time.format(datetime_fmt)

	if not wrf_work_dir:
		if tag != None:
			wrf_work_dir = f'{work_root}/{tag}/wrf'
		else:
			wrf_work_dir = f'{work_root}/wrf'

	if tag != None:
		obsproc_work_dir = f'{work_root}/{tag}/wrfda/obsproc'
	else:
		obsproc_work_dir = f'{work_root}/wrfda/obsproc'
	if not os.path.isdir(obsproc_work_dir): os.makedirs(obsproc_work_dir)
	old_dir = os.getcwd()
	os.chdir(obsproc_work_dir)

	cli.notice('Use builtin obserr.')
	run(f'ln -sf {wrfda_root}/var/obsproc/obserr.txt {obsproc_work_dir}')

	# Use d01 domain extent.
	if check_files([f'{wrf_work_dir}/wrfinput_d01_{start_time_str}']):
		ncfile       = Dataset(f'{wrf_work_dir}/wrfinput_d01_{start_time_str}', 'r')
		iproj        = ncfile.getncattr('MAP_PROJ')
		phic         = ncfile.getncattr('CEN_LAT')
		xlonc        = ncfile.getncattr('CEN_LON')
		moad_cen_lat = ncfile.getncattr('MOAD_CEN_LAT')
		standard_lon = ncfile.getncattr('STAND_LON')
		ncfile.close()
	else:
		iproj        = config['geogrid']['map_proj']
		phic         = config['geogrid']['ref_lat']
		xlonc        = config['geogrid']['ref_lon']
		moad_cen_lat = config['geogrid']['ref_lat']
		standard_lon = config['geogrid']['ref_lon']

	output_format = get_value(config, ['custom', 'obsproc', 'output_format'], default=2)
	time_window   = get_value(config, ['custom', 'wrfda', 'time_window'], default=360)

	if has_key(config, ('custom', 'da', 'type')):
		if config['custom']['da']['type'] == '3dvar':
			namelist_obsproc = f90nml.read(f'{wrfda_root}/var/obsproc/namelist.obsproc.3dvar.wrfvar-tut')
		else:
			cli.error('Currently, we only support 3DVar...')
	else:
		namelist_obsproc = f90nml.read(f'{wrfda_root}/var/obsproc/namelist.obsproc.3dvar.wrfvar-tut')

	namelist_obsproc['record1']['obs_gts_filename' ] = f'obs.gts.{start_time.format("YYYYMMDDHHmm")}'
	namelist_obsproc['record2']['time_window_min'  ] = start_time.subtract(minutes=time_window/2).format('YYYY-MM-DD_HH:mm:ss')
	namelist_obsproc['record2']['time_analysis'    ] = start_time.format('YYYY-MM-DD_HH:mm:ss')
	namelist_obsproc['record2']['time_window_max'  ] = start_time.add(minutes=time_window/2).format('YYYY-MM-DD_HH:mm:ss')
	namelist_obsproc['record3']['max_number_of_obs'] = 1200000
	namelist_obsproc['record7']['PHIC'             ] = phic
	namelist_obsproc['record7']['XLONC'            ] = xlonc
	namelist_obsproc['record7']['MOAD_CEN_LAT'     ] = moad_cen_lat
	namelist_obsproc['record7']['STANDARD_LON'     ] = standard_lon
	namelist_obsproc['record8']['NESTIX'           ] = config['geogrid']['e_sn']
	namelist_obsproc['record8']['NESTJX'           ] = config['geogrid']['e_we']
	namelist_obsproc['record8']['DIS'              ] = config['geogrid']['dx']
	namelist_obsproc['record9']['OUTPUT_OB_FORMAT' ] = output_format
	namelist_obsproc.write('./namelist.obsproc', force=True)

	cli.stage(f'Run obsproc.exe at {obsproc_work_dir} ...')
	expected_files = [f'obs_gts_{start_time.format("YYYY-MM-DD_HH:mm:ss")}.3DVAR']
	if not check_files(expected_files) or args.force:
		run('rm -f obs_gts_*')

		if littler_root == None: littler_root = args.littler_root
		if littler_root == None: cli.error('run_wrfda_obsproc: Argument littler_root is empty!')
		if has_key(config, ('custom', 'littler')):
			if 'dir_pattern' in config['custom']['littler'] and 'file_pattern' in config['custom']['littler']:
				dir_name = Template(config['custom']['littler']['dir_pattern']).render(time=start_time)
				if type(config['custom']['littler']['file_pattern']) == str:
					file_pattern = config['custom']['littler']['file_pattern']
				elif type(config['custom']['littler']['file_pattern']) == dict and has_key(config['custom']['littler']['file_pattern'], 'conv'):
					file_pattern = config['custom']['littler']['file_pattern']['conv']
				else:
					file_pattern = None
				if file_pattern != None:
					file_name = Template(file_pattern).render(time=start_time)
					littler_path = f'{littler_root}/{dir_name}/{file_name}'
				# Link Radar file.
				radar_path = None
				if type(config['custom']['littler']['file_pattern']) == dict and has_key(config['custom']['littler']['file_pattern'], 'radar'):
					file_pattern = config['custom']['littler']['file_pattern']['radar']
					file_name = Template(file_pattern).render(time=start_time)
					radar_path = f'{littler_root}/{dir_name}/{file_name}'
			else:
				cli.error('No dir_pattern and file_pattern in custom->littler section!')
		else:
			littler_path = f'{littler_root}/{start_time.format("YYYYMMDD")}/obs.gts.{start_time.format("YYYYMMDDHHmm")}'

		if os.path.exists(littler_path):
			if blacklist or whitelist:
				run(f'cp -f {littler_path} {obsproc_work_dir}/obs.gts.{start_time.format("YYYYMMDDHHmm")}')
			else:
				run(f'ln -sf {littler_path} {obsproc_work_dir}/obs.gts.{start_time.format("YYYYMMDDHHmm")}')
			littler_path = f'{obsproc_work_dir}/obs.gts.{start_time.format("YYYYMMDDHHmm")}'
		else:
			littler_path = None
		# Copy or link radar file.
		if radar_path != None and not os.path.exists(radar_path):
			cli.error(f'Radar file {radar_path} does not exist!')
		if radar_path != None:
			if blacklist or whitelist:
				run(f'cp -f {radar_path} {obsproc_work_dir}/obs.radar.{start_time.format("YYYYMMDDHHmm")}')
			else:
				run(f'ln -sf {radar_path} {obsproc_work_dir}/obs.radar.{start_time.format("YYYYMMDDHHmm")}')
			radar_path = f'{obsproc_work_dir}/obs.radar.{start_time.format("YYYYMMDDHHmm")}'

		if blacklist and not has_key(config, ('custom', 'blacklist')):
			cli.error(f'There is no "blacklist" in config object!')
		if whitelist and not has_key(config, ('custom', 'whitelist')):
			cli.error(f'There is no "whitelist" in config object!')

		if blacklist:
			for obs_type, sid_list in config['custom']['blacklist'].items():
				if obs_type == 'sound':
					tmp = 'FM-35 TEMP'
				elif obs_type == 'synop':
					tmp = 'FM-12 SYNOP'
				elif obs_type == 'amdar':
					tmp = 'FM-42 AMDAR'
				elif obs_type == 'profiler':
					tmp = 'FM-132 PROFLR'
				elif obs_type == 'gbmwr':
					tmp = 'FM-235 GBMWR'
				if sid_list == ['all']:
					if obs_type == 'radar':
						run(f"sed -i -E 's/^TOTAL NUMBER=[[:space:]]*([[:digit:]]+)/TOTAL NUMBER= 0/' {radar_path}")
					else:
						run(f"sed -i 's/^\\(.\\{{120\\}}{tmp}\\s\\+.\\{{139\\}}\\)F/\\1T/g' {obsproc_work_dir}/obs.gts.{start_time.format('YYYYMMDDHHmm')}")
				else:
					if type(sid_list) != list:
						sid_list = [sid_list]
					if obs_type == 'radar':
						for sid in sid_list:
							run(f"sed -i -E '/^RADAR[[:space:]]+{sid}/,/^RADAR/{{/^RADAR/!d;}}' {radar_path}")
							run(f"sed -i -E '/^RADAR[[:space:]]+{sid}/d' {radar_path}")
							n = int(re.search(r'\d+', open(radar_path, 'r').readline())[0])
							run(f"sed -i -E 's/^TOTAL NUMBER=[[:space:]]*([[:digit:]]+)/TOTAL NUMBER= {n-1}/' {radar_path}")
					else:
						for sid in sid_list:
							run(f"sed -i 's/^\\(.\\{{40\\}}{sid}\\s\\+.\\{{40\\}}{tmp}\\s\\+.\\{{139\\}}\\)F/\\1T/g' {obsproc_work_dir}/obs.gts.{start_time.format('YYYYMMDDHHmm')}")
		elif whitelist:
			# Turn off all observation!
			run(f"sed -i 's/^\\(.\\{{299\\}}\\)F/\\1T/g' {obsproc_work_dir}/obs.gts.{start_time.format('YYYYMMDDHHmm')}")
			for obs_type, sid_list in config['custom']['whitelist'].items():
				if obs_type == 'sound':
					tmp = 'FM-35 TEMP'
				elif obs_type == 'synop':
					tmp = 'FM-12 SYNOP'
				elif obs_type == 'amdar':
					tmp = 'FM-42 AMDAR'
				elif obs_type == 'profiler':
					tmp = 'FM-132 PROFLR'
				elif obs_type == 'gbmwr':
					tmp = 'FM-235 GBMWR'
				if sid_list == ['all']:
					if obs_type == 'radar':
						cli.error('Implement RADAR whitelist!')
					else:
						run(f"sed -i 's/^\\(.\\{{120\\}}{tmp}\\s\\+.\\{{139\\}}\\)T/\\1F/g' {obsproc_work_dir}/obs.gts.{start_time.format('YYYYMMDDHHmm')}")
				else:
					if type(sid_list) != list:
						sid_list = [sid_list]
					if obs_type == 'radar':
						for sid in sid_list:
							cli.error('Implement RADAR whitelist!')
					else:
						for sid in sid_list:
							run(f"sed -i 's/^\\(.\\{{40\\}}{sid}\\s\\+.\\{{40\\}}{tmp}\\s\\+.\\{{139\\}}\\)T/\\1F/g' {obsproc_work_dir}/obs.gts.{start_time.format('YYYYMMDDHHmm')}")

		if littler_path != None:
			submit_job(f'{wrfda_root}/var/obsproc/obsproc.exe', 1, config, args, wait=True)
			if not check_files(expected_files):
				cli.error(f'Failed!')
			run('ls -l obs_gts_*')
		if radar_path != None:
			run(f'ls -l {radar_path}')
		cli.notice('Succeeded.')
	else:
		cli.notice('File obs_gts_* already exist.')
	os.chdir(old_dir)
